import torch
import torch.nn as nn
import torch.nn.functional as F

# 定义一个Unimodal_GatedFusion类，继承自nn.Module
class Unimodal_GatedFusion(nn.Module):
    # 初始化函数，参数为隐藏层大小和数据集
    def __init__(self, hidden_size, dataset):
        super(Unimodal_GatedFusion, self).__init__()
        # 定义一个线性层，输入为隐藏层大小，输出为隐藏层大小，不使用偏置
        self.fc = nn.Linear(hidden_size, hidden_size, bias=False)
        # 如果数据集为MELD，则将线性层的权重设置为单位矩阵，并且不使用梯度
        # init.kaiming_uniform_(self.fc.weight, nonlinearity='relu')

        if dataset == 'MELD':
            # self.fc.weight.data.copy_(torch.eye(hidden_size, hidden_size))
            self.fc.weight.requires_grad = True

    # 定义前向传播函数，参数为a
    def forward(self, a):
        # 使用sigmoid函数对线性层的输出进行激活
        z = torch.sigmoid(self.fc(a))
        # 计算最终 representation
        final_rep = z * a
        # 返回最终 representation
        return final_rep